# 1、统一优化文件名
## 1.1 处理初始文件
Inital_image_preprocessing_Enable = True
# Inital_image_preprocessing_Enable = False

if Inital_image_preprocessing_Enable:
    import os
    folder = "/home/leo/Downloads/MVS pictures/28-all"
    picture_type = "28_"
    for i, filename in enumerate(os.listdir(folder)):
        timestaps = filename[6:]
        timestaps = timestaps[:-7]
        new_name = f"{picture_type + timestaps}.bmp"
        os.rename(os.path.join(folder, filename), os.path.join(folder, new_name))
        print(i, filename, new_name)

    folder = "/home/leo/Downloads/MVS pictures/32-all"
    picture_type = "32_"
    for i, filename in enumerate(os.listdir(folder)):
        timestaps = filename[6:]
        timestaps = timestaps[:-7]
        new_name = f"{picture_type + timestaps}.bmp"
        os.rename(os.path.join(folder, filename), os.path.join(folder, new_name))
        print(i, filename, new_name)

    folder = "/home/leo/Downloads/MVS pictures/36-all"
    picture_type = "36_"
    for i, filename in enumerate(os.listdir(folder)):
        timestaps = filename[6:]
        timestaps = timestaps[:-7]
        new_name = f"{picture_type + timestaps}.bmp"
        os.rename(os.path.join(folder, filename), os.path.join(folder, new_name))
        print(i, filename, new_name)

## 1.2 调整现有文件
# Current_Image_preprocessing_Enable = True
Current_Image_preprocessing_Enable = False

if Current_Image_preprocessing_Enable:
    import os
    folder_train = "C:\Temp\TF2.4 training\splited dataset\\train"
    folder_test = "C:\Temp\TF2.4 training\splited dataset\\test"

    def remane_files_in_folder(folder_path):
        for i, filename in enumerate(os.listdir(folder_path)):
            # filename_head = filename[:3]
            # filename_head = filename_head[:2]
            # filename_tail = filename[3:]
            # new_name = f"{filename_head + filename_tail}.bmp"
            new_name = filename[:-4]
            os.rename(os.path.join(folder_path, filename), os.path.join(folder_path, new_name))
            print(i, filename, new_name)

    remane_files_in_folder(folder_train)
    remane_files_in_folder(folder_test)

# 2、将每个类别的图片按比例随机分开（训练集、测试集）
# 将元数据随机打乱，并选取10%作为测试集（90%训练+验证，10%测试）
# Split_files_Enable = True
Split_files_Enable = False

input_folder = "C:\Temp\TF2.4 training\dataset_root"
output_folder = "C:\Temp\TF2.4 training\splited dataset"
# input_folder = "D:\Engineer_workshop\SW station\work\JDEC\project data\initial data"
# output_folder = "D:\Engineer_workshop\SW station\work\JDEC\project data\\testdata"
# input_folder = "D:\Engineer_workshop\SW station\work\JDEC\project data\initial data"
# output_folder = "D:\Engineer_workshop\SW station\work\JDEC\project data\initial data"

if Split_files_Enable:
    import splitfolders

    splitfolders.ratio(
        input_folder,
        output=output_folder,
        seed=42,
        ratio=(0.8, 0.1, 0.1),
        group_prefix=None
    )

# 3、遍历文件夹，生成标签文件
# Lable_files_Enable = True
Lable_files_Enable = False

# dataset_train = 'C:\Temp\TF2.4 training\splited dataset\\train'
# dataset_test = 'C:\Temp\TF2.4 training\splited dataset\\test'
# output_label_train = 'C:\Temp\TF2.4 training\splited dataset\labels-train.txt'
# output_label_test = 'C:\Temp\TF2.4 training\splited dataset\labels-test.txt'
dataset_train = 'C:\Temp\TF2.4 training\\test datasets\\train'
dataset_val = 'C:\Temp\TF2.4 training\\test datasets\\train'
dataset_test = 'C:\Temp\TF2.4 training\\test datasets\\test'
output_label_train = 'C:\Temp\TF2.4 training\\test datasets\labels-train.txt'
output_label_val = 'C:\Temp\TF2.4 training\\test datasets\labels-train.txt'
output_label_test = 'C:\Temp\TF2.4 training\\test datasets\labels-test.txt'

if Lable_files_Enable:
    import os

    def generate_img_labels(folder_path, output_path):
        files = os.listdir(folder_path)
        files.sort()
        with open(output_path, 'w') as f:
            for file in files:
                if not file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
                    continue

                if file[:2] == '28':
                    f.write(f"{file} {0}\n")
                else:
                    f.write(f"{file} {1}\n")

    generate_img_labels(dataset_train, output_label_train)
    generate_img_labels(dataset_test, output_label_test)